from backtrader import indicator
from backtrader.feed import DataClone
import pandas as pd
from pymongo import MongoClient
import datetime
import time
import math
import backtrader as bt
import backtrader.indicators as btind
import backtrader.feeds as btfeeds
import talib
from my_plot import *
from tqdm import tqdm
from py_mysql import *


class PandasData_more(bt.feeds.PandasData):
    pass
    # lines = ('m_RSI','h_RSI') # 要添加的线
    # # 设置 line 在数据源上的列位置
    # params=(
    #     ('m_RSI', 5),
    #     ('h_RSI', 7),
    #        ) 

# 创建策略继承bt.Strategy
class TestStrategy(bt.Strategy):
    params = (
        # 均线参数设置15天，15日均线
        ('maperiod', 15),
    )

    def log(self, txt, dt=None):
        # 记录策略的执行日志
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        # 保存收盘价的引用
        self.dataclose = self.datas[0].close
        # 跟踪挂单
        self.order = None
        # 买入价格和手续费
        self.buyprice = None
        self.buycomm = None
        #交易笔数
        self.bar_executed = 0


        # 23:00该时段平仓
        # self.close_time_arr = ['22:59:00']
        #  开始交易时间
        # self.status_time = '09:01:00'
        # 前十分钟趋势最大值
        self.trend_max = None
        # 前十分钟趋势最小值
        self.trend_min = None
        # 

        self.closeList = []


    def next(self):
        self.time = str(bt.num2date(self.lines.datetime[0]))
        self.position_size = self.getposition(self.datas[0]).size  #主期货交易量
        # 前十分钟趋势close数组
        self.h_m_s = (self.time).split(' ')[1]
        if hms_to_seconds(self.h_m_s) < 33060:
            self.closeList.append(self.dataclose[0])
        #  取出最大最小值，然后将数组清空
        elif hms_to_seconds(self.h_m_s) == 33060:
            self.trend_max = max(self.closeList)
            self.trend_min = min(self.closeList)
            self.closeList = []
        # 09:10至23:00 交易阶段
        elif hms_to_seconds(self.h_m_s) < 82800:
            self.max_min_num = (self.trend_max - self.trend_min)*0.618
            self.max_close = self.trend_max + self.max_min_num
            self.min_close = self.trend_min - self.max_min_num

            # if self.dataclose[0] > self.max_close or self.dataclose[0] < self.min_close:
            #     self.close()

            if self.dataclose[0] > self.trend_max:
                if self.position_size < 10:
                    self.buy(size=10)
            if self.dataclose[0] < self.trend_min:
                if self.position_size > -10:
                    self.sell(size=10)
            pass
        # 23:00 平仓
        elif hms_to_seconds(self.h_m_s) == 82800:
            self.trend_max = None
            self.trend_min = None
            self.close()
            
            print('该笔交易为平仓交易------size:',self.getposition(self.datas[0]).size)
            print('交易了:{}笔'.format(self.bar_executed))
            
        

     # 订单状态通知，买入卖出都是下单
    def notify_order(self, order):
        # 提交了/接受了,  买/卖订单什么都不做
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 检查一个订单是否完成
        # 注意:当资金不足时，broker会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm

                print({'状态':'买入','date': self.time, 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
            elif order.issell():
                print({'状态':'卖出','date': self.time, 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
            self.bar_executed += 1
            

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            print('订单取消/保证金不足/拒绝')
        # 其他状态记录为：无法挂单
        self.order = None


    # 交易状态通知，一买一卖算交易（交易净利润）
    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        


def main(codeArr, startDate, endDate, startcash=10000000, com=0.00002, qts=13):

    # 创建主控制器
    cerebro = bt.Cerebro()
    # 导入策略参数寻优
    cerebro.addstrategy(TestStrategy)

    # 获取数据
    query_db = Mysql_search()
    df = query_db.get_one(codeArr,startDate,endDate)
    for item in df:
        df2 = df[item]
        df2.index = pd.to_datetime( df2.index, utc= True)
        print(df2)
        data = PandasData_more(
            dataname = df2,
            timeframe=bt.TimeFrame.Minutes)
        cerebro.adddata(data)

    # broker设置资金、手续费
    cerebro.broker.setcash(startcash)
    cerebro.broker.setcommission(commission=com)
    # 设置买入设置，策略，数量
    cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
    # 以发出信号当日收盘价成交
    cerebro.broker.set_coc(True)
    
    print('期初总资金: %.2f' % cerebro.broker.getvalue())

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name= '_TimeReturn')

    result = cerebro.run()

    print('期末总资金: %.2f' % cerebro.broker.getvalue())

    custom_plot(result)
    # cerebro.plot()


def hms_to_seconds(t):
    h, m, s = [int(i) for i in t.split(':')]
    return 3600*h + 60*m + s
# print(hms_to_seconds('23:00:00'))

if __name__ == '__main__':
    startDate = '2021-08-01'  #开始时间
    endDate = '2021-11-09'   #结束时间
    codeArr = ['RB']   #品种数组
    
    main(codeArr,startDate,endDate)







